package com.sui.bigdata.sml.web.misc.azkaban;

import com.alibaba.fastjson.JSON;
import com.feidee.fdcommon.azkaban.constant.SparkConstant;
import com.feidee.fdcommon.azkaban.model.SparkJobParam;
import com.feidee.fdcommon.configuration.CustomConfiguration;
import com.feidee.fdcommon.util.ToolUtils;
import com.sui.bigdata.sml.web.dto.ConnectionDTO;
import com.sui.bigdata.sml.web.dto.NodeDTO;
import com.sui.bigdata.sml.web.repository.model.AzProject;
import org.apache.commons.text.StringEscapeUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author songhaicheng
 * @date 2019/9/26 14:08
 * @description Azkaban 任务处理类
 * @reviewer
 */
public class AzkabanHelper {

    public static String getAZUsername() {
        return CustomConfiguration.getString(SparkConstant.SML_AZ_USERNAME, "admin");
    }

    /**
     * 初始化要提交的 project
     *
     * @return
     */
    public static AzProject initProject() {
        return new AzProject()
                .setCron("")
                .setUsername(getAZUsername())
                .setEmails("")
                .setName("SML_" + System.currentTimeMillis())
                .setFailureAction("02");
    }

    /**
     * 给一个 az job 命名
     *
     * @param jobId 节点名字
     * @return
     */
    private static SparkJobParam initJob(String jobId) {
        SparkJobParam jobParam = new SparkJobParam();
        jobParam.setType("java");
        jobParam.setRetries(0);
        jobParam.setJobId(jobId);
        return jobParam;
    }

    /**
     * 把前端的节点信息转成 az job 实体
     *
     * @param node         前端节点信息
     * @param algClassName 运行类路径参数
     * @param nodeGroupId  运行队列 id
     * @return
     */
    public static SparkJobParam createSparkJob(Node node, String algClassName, Long nodeGroupId) {
        SparkJobParam jobParam = initJob(node.val.getName());
        jobParam.setAlgClassName(algClassName);
        if (ToolUtils.isNotNull(node.val.getParameters())) {
            Map<String, Object> algParams = JSON.parseObject(node.val.getParameters());
            // 对于运行参数，不解析转义字符，防止传到 algorithm 组件出错
            algParams.forEach((k, v) -> {
                if (v instanceof String) {
                    algParams.put(k, StringEscapeUtils.escapeJava(v.toString()));
                }
            });
            algParams.put(SparkConstant.SML_EXECUTION_LOG_DIR,
                    CustomConfiguration.getString(SparkConstant.SML_EXECUTION_LOG_DIR));
            algParams.put(SparkConstant.SML_NODE_ID, node.val.getId());
            algParams.put(SparkConstant.SML_NODE_GROUP_ID, nodeGroupId);
            // 设置节点输入信息
            if (ToolUtils.isNotNull(node.val.getInput())) {
                algParams.put(SparkConstant.INPUT_PATH, node.val.getInput());
            }
            jobParam.setAlgParams(algParams);
        }
        // 构建依赖关系
        jobParam.setDependencies(node.parents.stream().map(p -> p.val.getName()).toArray(String[]::new));
        // Advanced params
        if (ToolUtils.isNotNull(node.val.getRunningConf())) {
            Map<String, Object> conf = JSON.parseObject(node.val.getRunningConf());
            jobParam.setDriverMemory((String) conf.getOrDefault(SparkConstant.SPARK_DRIVER_MEMORY, ""));
            jobParam.setNumExecutors((Integer) conf.getOrDefault(SparkConstant.SPARK_NUM_EXECUTORS, 0));
            jobParam.setExecutorMemory((String) conf.getOrDefault(SparkConstant.SPARK_EXECUTOR_MEMORY, " "));
            jobParam.setExecutorCores((Integer) conf.getOrDefault(SparkConstant.SPARK_EXECUTOR_CORES, 0));
            jobParam.setConf((String) conf.getOrDefault(SparkConstant.SPARK_CONF, ""));
        }
        return jobParam;
    }

    /**
     * 在 flow 末尾加个 endJob，以实现 az 的任务并行化
     *
     * @param experimentName 所属实验名
     * @param nodes          解析出来的运行节点
     * @return
     */
    public static SparkJobParam appendEndJob(String experimentName, List<Node> nodes) {
        SparkJobParam endJob = initJob("FLOW_" + experimentName);
        endJob.setType("cmd");
        endJob.setAlgClassName("end_job");
        endJob.setDependencies(
                nodes.parallelStream()
                        .filter(n -> n.children.size() == 0)
                        .map(n -> n.val.getName())
                        .toArray(String[]::new)
        );
        return endJob;
    }

    /**
     * 构建运行图
     *
     * @param id2Node     所在实验的所有节点与其 id 映射表
     * @param connections 连线
     * @param entry       在前端界面点击的目标节点（为空时，代表运行整个实验）
     * @return 所有会跑的节点
     */
    public static List<Node> initGraph(Map<Long, NodeDTO> id2Node, List<ConnectionDTO> connections, NodeDTO entry)
            throws RuntimeException {
        Map<Long, Node> idToGraphNode = new HashMap<>();
        Queue<Long> q = new LinkedList<>();
        if (entry != null) {
            q.offer(entry.getId());

        } else {
            // 找到所有根节点，并加入队列中
            List<Long> tails = connections.parallelStream()
                    .map(ConnectionDTO::getTailNid)
                    .collect(Collectors.toList());
            id2Node.values().parallelStream()
                    .map(NodeDTO::getId)
                    .filter(nid -> !tails.contains(nid))
                    .forEach(q::offer);
        }
        while (!q.isEmpty()) {
            Long nid = q.poll();
            Node node = idToGraphNode.getOrDefault(nid, new Node(id2Node.get(nid)));
            List<Long> childrenIds = connections.parallelStream()
                    .filter(conn -> conn.getHeadNid().equals(nid))
                    .map(ConnectionDTO::getTailNid)
                    .collect(Collectors.toList());
            childrenIds.forEach(id -> {
                Node child = idToGraphNode.getOrDefault(id, new Node(id2Node.get(id)));
                child.parents.add(node);
                node.children.add(child);
                idToGraphNode.put(child.val.getId(), child);
                q.offer(child.val.getId());
            });
            idToGraphNode.put(node.val.getId(), node);
        }
        // 如果有节点未设置好必填参数，直接报错
        idToGraphNode.values().forEach(node -> {
            if (node.val.getStatus() == 0) {
                throw new RuntimeException("组件【" + node.val.getName() + "】必填参数未设置完整");
            }
        });
        return new ArrayList<>(idToGraphNode.values());
    }

    /**
     * 初始化运行节点信息（状态（涉及到前端图标刷新），输入、输出）
     *
     * @param nodes 目标节点们
     */
    public static void initNodes(List<Node> nodes) {
        nodes.parallelStream().forEach(n -> {
            // 如果是顶部节点，保留原始输入数据信息
            if (n.parents.size() == 0) {
                n.val.initStatus();
                n.val.initOutput();
            } else {
                n.val.init();
            }
        });
    }

    /**
     * 抽象化的节点
     */
    public static class Node {
        /**
         * 前端节点实体
         */
        public NodeDTO val;
        /**
         * 父节点
         */
        public List<Node> parents = new ArrayList<>();
        /**
         * 子节点
         */
        public List<Node> children = new ArrayList<>();

        public Node(NodeDTO val) {
            this.val = val;
        }
    }

    public enum RunningType {
        /**
         * 单个
         */
        SINGLE("single"),
        /**
         * 从此开始
         */
        FROM("from"),
        /**
         * 运行到此
         */
        TO("to"),
        /**
         * 运行所有
         */
        ALL("all");

        private String value;

        RunningType(String value) {
            this.value = value;
        }

        public String getValue() {
            return this.value;
        }

        public static RunningType getRunningType(String candidate) {
            for (RunningType type : values()) {
                if (type.getValue().equals(candidate.toLowerCase())) {
                    return type;
                }
            }
            throw new RuntimeException("无效的运行方式：" + candidate);
        }

    }

}
